# groundingdino_server.py
import warnings
warnings.filterwarnings("ignore", message="Importing from timm.models.layers is deprecated")
warnings.filterwarnings("ignore", message="torch.meshgrid: in an upcoming release")
warnings.filterwarnings("ignore", message="The `device` argument is deprecated")

from groundingdino.util.inference import load_model, load_image, predict, annotate
import cv2
import base64
import numpy as np
import logging
from datetime import datetime
import argparse
import sys
from flask import Flask, request, jsonify
from flask_cors import CORS
import threading
import time
import torch
from PIL import Image
import torchvision.transforms as T

# 解析命令行参数
parser = argparse.ArgumentParser()
parser.add_argument('--log-file', help='日志文件路径')
args = parser.parse_args()

# 配置日志
if args.log_file:
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(args.log_file, encoding='utf-8'),
            logging.StreamHandler(sys.stdout)
        ]
    )
else:
    logging.basicConfig(level=logging.INFO)

logger = logging.getLogger(__name__)

# 全局变量用于存储模型
model = None

# 默认参数
BOX_THRESHOLD = 0.35
TEXT_THRESHOLD = 0.25

# 图像预处理变换
transform = T.Compose([
    T.ToTensor(),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

def initialize_model():
    """初始化并加载模型"""
    global model
    logger.info("开始加载GroundingDINO模型...")
    model = load_model("groundingdino/config/GroundingDINO_SwinT_OGC.py", "weights/groundingdino_swint_ogc.pth")
    logger.info("GroundingDINO模型加载成功")

def base64_to_image(base64_string):
    """将base64字符串转换为OpenCV图像"""
    # 移除base64字符串中的前缀（如果存在）
    if base64_string.startswith('data:image'):
        base64_string = base64_string.split(',')[1]
    
    # 解码base64字符串
    image_data = base64.b64decode(base64_string)
    
    # 转换为numpy数组
    np_array = np.frombuffer(image_data, np.uint8)
    
    # 解码为OpenCV图像
    image = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
    
    return image

def image_to_base64(image):
    """将OpenCV图像转换为base64字符串"""
    _, buffer = cv2.imencode('.jpg', image, [int(cv2.IMWRITE_JPEG_QUALITY), 80])
    jpg_as_text = base64.b64encode(buffer).decode('utf-8')
    return jpg_as_text

def detect_objects_with_text(base64_image, text_prompt):
    """核心目标检测方法"""
    start_time = datetime.now()
    
    # 检查模型是否已加载
    if model is None:
        logger.error("模型未加载")
        return {'error': '模型未加载'}, 500
        
    # 将base64图像转换为OpenCV格式
    image = base64_to_image(base64_image)
    
    if image is None:
        logger.warning("无法解码图像")
        return {'error': '无法解码图像'}, 400
    
    logger.info(f"开始进行GroundingDINO检测，提示文本: {text_prompt}")
    
    # 转换图像格式以适应GroundingDINO
    image_source = image
    # 转换为RGB
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # 将numpy数组转换为PIL图像，再转换为张量
    pil_image = Image.fromarray(image_rgb)
    image_tensor = transform(pil_image)
    
    # 进行目标检测
    boxes, logits, phrases = predict(
        model=model,
        image=image_tensor,
        caption=text_prompt,
        box_threshold=BOX_THRESHOLD,
        text_threshold=TEXT_THRESHOLD
    )
    
    # 绘制检测框
    annotated_frame = annotate(
        image_source=image_source, 
        boxes=boxes, 
        logits=logits, 
        phrases=phrases
    )
    
    # 提取检测信息
    detection_data = []
    for i in range(len(boxes)):
        box = boxes[i]
        logit = float(logits[i]) if logits is not None else 0.0
        phrase = phrases[i] if phrases is not None else "Unknown"
        
        # 将归一化的坐标转换为像素坐标
        h, w = image.shape[:2]
        x1 = int(box[0] * w)
        y1 = int(box[1] * h)
        x2 = int(box[2] * w)
        y2 = int(box[3] * h)
        
        detection_info = {
            'class_name': phrase,
            'confidence': logit,
            'bbox': [x1, y1, x2, y2]
        }
        detection_data.append(detection_info)
    
    logger.info(f"检测完成，发现 {len(detection_data)} 个目标")
    
    # 将处理后的图像转换为base64
    processed_image_base64 = image_to_base64(annotated_frame)
    
    if processed_image_base64 is None:
        logger.error("图像处理失败")
        return {'error': '图像处理失败'}, 500
    
    processing_time = (datetime.now() - start_time).total_seconds()
    
    # 返回结果
    response_data = {
        'processed_image': processed_image_base64,
        'detections': detection_data,
        'status': 'success',
        'processing_time': processing_time,
        'detection_count': len(detection_data),
        'prompt': text_prompt
    }
    
    logger.info(f"检测完成: 发现 {len(detection_data)} 个目标，处理时间: {processing_time:.2f}秒")
    return response_data, 200

def groundingdino_health_check():
    """GroundingDINO服务健康检查"""
    logger.info("健康检查请求")
    model_status = "loaded" if model is not None else "not loaded"
    return {
        'status': 'healthy', 
        'message': f'GroundingDINO服务正在运行，模型状态: {model_status}',
        'timestamp': datetime.now().isoformat()
    }

# 创建Flask应用
app = Flask(__name__)
CORS(app)

@app.route('/predict', methods=['POST'])
def predict_api():
    """接收base64图像和文本提示并返回检测结果"""
    logger.info("收到预测请求")
    data = request.get_json()
   
    if not data or 'image' not in data:
        logger.warning("缺少图像数据")
        return jsonify({'error': '缺少图像数据'}), 400
        
    if 'prompt' not in data:
        logger.warning("缺少文本提示数据")
        return jsonify({'error': '缺少文本提示数据'}), 400
    
    logger.debug("开始处理图像数据和文本提示")
    result, status_code = detect_objects_with_text(data['image'], data['prompt'])
    logger.info(f"预测请求处理完成，状态码: {status_code}")
    return jsonify(result), status_code

@app.route('/health', methods=['GET'])
def health_check_api():
    """健康检查接口"""
    logger.info("收到健康检查请求")
    return jsonify(groundingdino_health_check())

# 添加一个根路径路由，用于测试服务器是否正常运行
@app.route('/', methods=['GET'])
def index():
    return jsonify({"message": "GroundingDINO服务正在运行"}), 200

def start_server():
    """启动Flask服务器"""
    logger.info("启动GroundingDINO Flask服务器...")
    logger.info("监听端口: 5005")
    app.run(host='0.0.0.0', port=5005, debug=False)

# 在文件最后启动服务器
if __name__ == "__main__":
    logger.info("启动GroundingDINO预测服务...")
    # 立即加载模型
    initialize_model()
    
    # 启动Flask服务器线程
    server_thread = threading.Thread(target=start_server, daemon=True)
    server_thread.start()
    
    # 保持主线程运行
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        logger.info("服务器已停止")